{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import gym\n",
    "import matplotlib.pyplot as plt\n",
    "import copy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "WINDOW_SIZE = 128; NUM_FEATURES = 6; NUM_ACTIONS = 3; \n",
    "BATCH_SIZE = 128; LR = 0.01; EPISILO = 0.9\n",
    "MEMORY_CAPACITY = 2000; Q_NETWORK_ITERATION = 100; GAMMA = 0.99"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df = pd.read_csv('adj/1.csv')\n",
    "\n",
    "df['high'] = df['high']*df['factor']\n",
    "df['low'] = df['low']*df['factor']\n",
    "df['open'] = df['open']*df['factor']\n",
    "df['close'] = df['close']*df['factor']\n",
    "\n",
    "df_high = df.groupby('date')[['high']].max().reset_index(drop=False)\n",
    "df_new = df_high\n",
    "\n",
    "df_low = df.groupby('date')[['low']].min().reset_index(drop=False)\n",
    "df_new = pd.merge(df_new,df_low,how='left',on=['date'])\n",
    "\n",
    "list_open = list(df.groupby('date')['date_time'].min())\n",
    "df_open = df[df['date_time'].isin(list_open)][['date','open']].reset_index(drop=True)\n",
    "df_new = pd.merge(df_new,df_open,how='left',on=['date'])\n",
    "\n",
    "list_close = list(df.groupby('date')['date_time'].max())\n",
    "df_close = df[df['date_time'].isin(list_close)][['date','close']].reset_index(drop=True)\n",
    "df_new = pd.merge(df_new,df_close,how='left',on=['date'])\n",
    "\n",
    "df_volume = df.groupby('date')[['volume']].sum().reset_index(drop=False)\n",
    "df_new = pd.merge(df_new,df_volume,how='left',on=['date'])\n",
    "\n",
    "df_amount = df.groupby('date')[['amount']].sum().reset_index(drop=False)\n",
    "df_new = pd.merge(df_new,df_amount,how='left',on=['date'])\n",
    "\n",
    "df_new['date_time'] = df_new['date']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "dqn = DQN()\n",
    "env = Env(data=df_new, window_size=WINDOW_SIZE, tranct=0.0002, position_limit=5)\n",
    "episodes = 400\n",
    "print(\"Collecting Experience....\")\n",
    "reward_list = []\n",
    "plt.ion()\n",
    "fig, ax = plt.subplots()\n",
    "for i in range(episodes):\n",
    "    state = env.reset()\n",
    "    ep_reward = 0\n",
    "    while True:\n",
    "        action = dqn.choose_action(state)\n",
    "        x = env.act(action)\n",
    "        next_state = x[0]; reward = x[1]; done = x[2]\n",
    "        dqn.store_transition(state, action, reward, next_state)\n",
    "        ep_reward += reward\n",
    "        if dqn.memory_counter >= MEMORY_CAPACITY:\n",
    "            dqn.learn()\n",
    "            if done:\n",
    "                print(\"episode: {} , the episode reward is {}\".format(i, round(ep_reward, 3)))\n",
    "        if done:\n",
    "            break\n",
    "        state = next_state\n",
    "    r = copy.copy(reward)\n",
    "    reward_list.append(r)\n",
    "    ax.set_xlim(0,300)\n",
    "    #ax.cla()\n",
    "    ax.plot(reward_list, 'g-', label='total_loss')\n",
    "    plt.pause(0.001)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Net(nn.Module):\n",
    "    def __init__(self,num_features,num_actions):\n",
    "        super(Net, self).__init__()\n",
    "        self.gru = nn.GRU(input_size=num_features,hidden_size=50, num_layers=2,batch_first=True)\n",
    "        self.fc = nn.Linear(50,30)\n",
    "        self.out = nn.Linear(30,num_actions)\n",
    "        self.softmax = nn.Softmax(dim=1)\n",
    "\n",
    "    def forward(self,x):\n",
    "        x = self.gru(x)\n",
    "        x = x[0][:,-1,:]\n",
    "        x = self.fc(x)\n",
    "        x = F.relu(x)\n",
    "        action_prob = self.softmax(self.out(x))\n",
    "        return action_prob\n",
    "\n",
    "class DQN():\n",
    "    \"\"\"docstring for DQN\"\"\"\n",
    "    def __init__(self):\n",
    "        super(DQN, self).__init__()\n",
    "        self.eval_net, self.target_net = Net(NUM_FEATURES,NUM_ACTIONS), Net(NUM_FEATURES,NUM_ACTIONS)\n",
    "\n",
    "        self.learn_step_counter = 0\n",
    "        self.memory_counter = 0\n",
    "        self.memory_state = np.zeros((MEMORY_CAPACITY, WINDOW_SIZE ,NUM_FEATURES * 2))\n",
    "        self.memory_actionvalue = np.zeros((MEMORY_CAPACITY,1 , 2))\n",
    "        # When we store the memory, we put the state, action, reward and next_state in the memory\n",
    "        # here reward and action is a number, state is a ndarray\n",
    "        self.optimizer = torch.optim.Adam(self.eval_net.parameters(), lr=LR)\n",
    "        self.loss_func = nn.MSELoss()\n",
    "\n",
    "    def choose_action(self, state):\n",
    "        state = torch.FloatTensor(np.expand_dims(state,axis=0)) # get a 1D array\n",
    "        if np.random.randn() <= EPISILO:# greedy policy\n",
    "            action_value = self.eval_net.forward(state)\n",
    "            action = action_value.argmax().item()\n",
    "        else: # random policy\n",
    "            action = np.random.randint(0,NUM_ACTIONS)\n",
    "        return action\n",
    "\n",
    "\n",
    "    def store_transition(self, state, action, reward, next_state):\n",
    "        transition_state = np.concatenate([state,next_state],axis=1)\n",
    "        transition_actionvalue = np.array([action,reward])\n",
    "        index = self.memory_counter % MEMORY_CAPACITY\n",
    "        self.memory_state[index,:,:] = transition_state\n",
    "        self.memory_actionvalue[index,:] = transition_actionvalue\n",
    "        self.memory_counter += 1\n",
    "\n",
    "\n",
    "    def learn(self):\n",
    "\n",
    "        #update the parameters\n",
    "        if self.learn_step_counter % Q_NETWORK_ITERATION ==0:\n",
    "            self.target_net.load_state_dict(self.eval_net.state_dict())\n",
    "        self.learn_step_counter+=1\n",
    "\n",
    "        #sample batch from memory\n",
    "        sample_index = np.random.choice(MEMORY_CAPACITY, BATCH_SIZE)\n",
    "        batch_memory_state = self.memory_state[sample_index]\n",
    "        batch_memory_actionvalue = self.memory_actionvalue[sample_index]\n",
    "        batch_state = torch.FloatTensor(batch_memory_state[:,:, :NUM_FEATURES])\n",
    "        batch_action = torch.LongTensor(batch_memory_actionvalue[:,:,0].astype(int))\n",
    "        batch_reward = torch.FloatTensor(batch_memory_actionvalue[:,:,1])\n",
    "        batch_next_state = torch.FloatTensor(batch_memory_state[:,:,-NUM_FEATURES:])\n",
    "\n",
    "        #q_eval\n",
    "        q_eval = self.eval_net(batch_state).gather(1, batch_action)\n",
    "        q_next = self.target_net(batch_next_state).detach()\n",
    "        q_target = batch_reward + GAMMA * q_next.max(1)[0].view(BATCH_SIZE, 1)\n",
    "        loss = self.loss_func(q_eval, q_target)\n",
    "\n",
    "        self.optimizer.zero_grad()\n",
    "        loss.backward()\n",
    "        self.optimizer.step()\n",
    "\n",
    "class Env():\n",
    "    def __init__(self, data, window_size, tranct, position_limit):\n",
    "        \n",
    "        if len(data)<window_size * 2:\n",
    "            raise Exception('data length < window_size')\n",
    "            \n",
    "        self.data = data\n",
    "        self.window_size = window_size\n",
    "        self.tranct = tranct\n",
    "        self.position_limit = position_limit\n",
    "            \n",
    "        self.info = pd.DataFrame({'date_time':[],'action':[],'price':[],'position':[]})\n",
    "        self.positions = []\n",
    "        self.t = window_size\n",
    "        self.state = None\n",
    "        self.reward = None\n",
    "        self.done = False\n",
    "        \n",
    "    def normalization(self,data):\n",
    "        mu = np.mean(data, axis=0)\n",
    "        sigma = np.std(data, axis=0)\n",
    "        return (data - mu)/sigma\n",
    "    \n",
    "    def reset(self):\n",
    "        self.t = self.window_size\n",
    "        self.info = pd.DataFrame({'date_time':[],'action':[],'price':[],'position':[]})\n",
    "        temp_df = pd.DataFrame({'date_time':[self.data.loc[self.t,'date_time']],'action':['initial'],\n",
    "                                    'price':[self.data.loc[self.t,'close']],'position':[0]})\n",
    "        self.info = self.info.append(temp_df).reset_index(drop=True)\n",
    "        self.state = self.normalization(self.data.loc[(self.t-self.window_size+1):self.t,['open','high','low','close','volume','amount']].values)\n",
    "        self.reward = 0\n",
    "        self.done = False\n",
    "        return self.state\n",
    "        \n",
    "    def act(self,action):  #buy,sell,sit\n",
    "        if (self.t+2) ==len(self.data):\n",
    "            self.t += 1\n",
    "            self.done = True\n",
    "            self.state = self.normalization(self.data.loc[(self.t-self.window_size+1):self.t,['open','high','low','close','volume','amount']].values)\n",
    "            self.reward = self.data.loc[self.t,'close']*(1-self.tranct)*len(self.positions) - sum(self.positions)\n",
    "            print('final reward is:',self.reward)\n",
    "            return (self.state,self.reward,self.done)\n",
    "        elif action==2:\n",
    "            self.t += 1\n",
    "            self.state = self.normalization(self.data.loc[(self.t-self.window_size+1):self.t,['open','high','low','close','volume','amount']].values)\n",
    "            self.reward = 0\n",
    "            return (self.state,self.reward,self.done)\n",
    "        elif action==0:\n",
    "            if (len(self.positions)+1)<self.position_limit:\n",
    "                temp_df = pd.DataFrame({'date_time':[self.data.loc[self.t,'date_time']],'action':['buy'],\n",
    "                                        'price':[self.data.loc[self.t,'close']*(1+self.tranct)],\n",
    "                                        'position':[(len(self.positions)+1)]})\n",
    "                self.info = self.info.append(temp_df).reset_index(drop=True)\n",
    "                self.positions.append(self.data.loc[self.t,'close']*(1+self.tranct))\n",
    "                print('buy at:',self.data.loc[self.t,'close']*(1+self.tranct))\n",
    "            else:\n",
    "                pass\n",
    "            self.t += 1\n",
    "            self.state = self.normalization(self.data.loc[(self.t-self.window_size+1):self.t,['open','high','low','close','volume','amount']].values)\n",
    "            self.reward = 0\n",
    "            return (self.state,self.reward,self.done)\n",
    "        elif action==1:\n",
    "            if len(self.positions)>0:\n",
    "                temp_df = pd.DataFrame({'date_time':[self.data.loc[self.t,'date_time']],'action':['sell'],\n",
    "                                        'price':[self.data.loc[self.t,'close']*(1-self.tranct)],\n",
    "                                        'position':[len(self.positions)-1]})\n",
    "                self.info = self.info.append(temp_df).reset_index(drop=True)\n",
    "                self.reward = self.data.loc[self.t,'close']*(1-self.tranct) - self.positions.pop(0)\n",
    "                print('sell at:',self.data.loc[self.t,'close']*(1-self.tranct),'reward is:',self.reward)\n",
    "            else:\n",
    "                self.reward = 0\n",
    "            self.t += 1\n",
    "            self.state = self.normalization(self.data.loc[(self.t-self.window_size+1):self.t,['open','high','low','close','volume','amount']].values)\n",
    "            return  (self.state,self.reward,self.done)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.8"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
